package ws

import (
	"github.com/gorilla/websocket"
	"github.com/morgine/log"
	"io"
	"sync"
	"time"
)

const (
	readDeadLine  = 10 * time.Second
	writeDeadLine = 10 * time.Second
)

type Client struct {
	conn        *websocket.Conn
	url         string
	done        chan struct{}
	handler     func(message []byte)
	onConnected func(w io.Writer)
	runOnce     sync.Once
	writeMutex  sync.Mutex
}

func NewClient(url string, onConnected func(w io.Writer), handler func(message []byte)) *Client {
	done := make(chan struct{})
	return &Client{url: url, handler: handler, done: done, onConnected: onConnected}
}

func (c *Client) Run() {
	c.runOnce.Do(func() {
		go c.autoConnect()
	})
}

func (c *Client) autoConnect() {
	c.connect()
	ticker := time.NewTicker(5 * time.Second)
	for {
		select {
		case <-c.done:
			return
		case <-ticker.C:
			log.Info.Printf("ws: retry connect %s\n", c.url)
			c.connect()
		}
	}
}

func (c *Client) connect() {
	conn, _, err := websocket.DefaultDialer.Dial(c.url, nil)
	if err != nil {
		log.Error.Printf("ws: connect %s: %s\n", c.url, err)
	} else {
		c.conn = conn
		log.Info.Printf("ws: connected %s\n", c.url)
		if c.onConnected != nil {
			c.onConnected(c)
		}
		c.read()
	}
}

func (c *Client) Write(data []byte) (int, error) {
	c.writeMutex.Lock()
	defer c.writeMutex.Unlock()
	err := c.conn.SetWriteDeadline(time.Now().Add(writeDeadLine))
	if err != nil {
		return 0, err
	}
	//log.Info.Printf("write: %s\n", data)
	return len(data), c.conn.WriteMessage(websocket.TextMessage, data)
}

// 开始监听消息
func (c *Client) read() {
	for {
		select {
		case <-c.done:
			return
		default:
			err := c.conn.SetReadDeadline(time.Now().Add(readDeadLine))
			if err != nil {
				log.Error.Printf("ws: set read dead line: %s\n", err)
				return
			}
			t, message, err := c.conn.ReadMessage()
			if err != nil {
				log.Error.Printf("ws: read message: %s\n", err)
				return
			}
			// 币安会主动发送 ping 针，返回 pong 针
			if t == websocket.PingMessage {
				err = c.conn.WriteControl(websocket.PongMessage, []byte("pong"), time.Now().Add(time.Second))
				if err != nil {
					log.Error.Printf("ws: response pong control: %s\n", err)
					return
				}
			} else {
				// 火币会发送一个 ping 对象，需要返回一个 ping 对象
				// OKEx 需要主动发送 一个 ping 字符串
				c.handler(message)
			}
		}
	}
}

func (c *Client) Close() error {
	//c.cancelTicker() // cancel ticker first
	c.done <- struct{}{}
	return c.conn.Close()
}
